"""
生成frame_list中的csv文件
"""
import os
import csv

def generate_csv(folder_path, train_csv_path, val_csv_path):
    # 创建训练和验证集 CSV 文件
    with open(train_csv_path, 'w', newline='') as train_csvfile, open(val_csv_path, 'w', newline='') as val_csvfile:
        train_writer = csv.writer(train_csvfile)
        val_writer = csv.writer(val_csvfile)

        # 写入合并后的表头，列之间用空格分隔
        train_writer.writerow(['original_video_id video_id frame_id path labels'])
        val_writer.writerow(['original_video_id video_id frame_id path labels'])

        # 获取视频文件夹并按数字顺序排序
        video_folders = sorted(os.listdir(folder_path), key=lambda x: int(x))

        # 遍历视频文件夹并按新规则划分
        for video_id, video_folder in enumerate(video_folders, start=1):
            video_folder_path = os.path.join(folder_path, video_folder)
            if os.path.isdir(video_folder_path):
                # 确定当前视频是否分配到训练或验证集
                is_validation = (video_id - 1) % 5 == 4  # 每组的第五个和第十个视频分配给验证集
                writer = val_writer if is_validation else train_writer

                # 获取帧文件并按数字顺序排序
                frame_files = sorted(os.listdir(video_folder_path), key=lambda x: int(x.split('.')[0]))

                # 遍历视频文件夹中的所有帧
                for frame_id, frame_file in enumerate(frame_files):
                    if frame_file.lower().endswith(('.jpg', '.png', '.jpeg')):  # 仅处理图片文件
                        frame_path = os.path.join(video_folder, frame_file)
                        frame_path = frame_path.replace('\\', '/')
                        # 合并列为一列，并写入 CSV 文件
                        row = f"{video_folder} {video_id - 1} {frame_id} {frame_path} \"\""
                        writer.writerow([row])

# 设置主文件夹路径
main_folder_path = r'D:\file\postgrad\experiment\bird_ava_dataset\videos_rawframes'  # 替换为你的主文件夹路径

# 设置输出文件路径
train_csv_file = r'D:\file\postgrad\experiment\bird_ava_dataset\frame_lists\train.csv'  # 训练集 CSV 文件路径
val_csv_file = r'D:\file\postgrad\experiment\bird_ava_dataset\frame_lists\val.csv'  # 验证集 CSV 文件路径

# 生成训练集和验证集 CSV 文件
generate_csv(main_folder_path, train_csv_file, val_csv_file)


